import numpy as np
import torch
from sklearn.manifold import TSNE
from torch.utils.data import DataLoader

from MyModel.models.MyModel import MyModel
from MyModel.models.MyModelWithoutLayer2 import MyModelWithoutLayer2
from MyModel.models.MyModleConfig import MyModelConfig
from dataset import TaoBaoDataset

from config import *
import os
import warnings
import matplotlib.pyplot as plt

os.environ['CUDA_LAUNCH_BLOCKING'] = '1'

warnings.filterwarnings('ignore')

def prepare_model_data(model_name, checkpoints_path,chose_loader):
    train_target_item_file = './data/taobao/final_train_data.csv'
    test_target_item_file = './data/taobao/final_test_data.csv'
    user_behavior_file = f'./data/taobao/active_user_actions.csv'
    n_users = 100
    print(model_name)

    if chose_loader=='test':
        train_test_data = TaoBaoDataset(test_target_item_file, user_behavior_file, mode='test', n_users=n_users, q_num=q_num)
        loader = DataLoader(train_test_data, batch_size=1, shuffle=False)
    elif chose_loader=='train':
        train_test_data = TaoBaoDataset(train_target_item_file, user_behavior_file, mode='train', n_users=n_users,q_num=q_num)
        loader = DataLoader(train_test_data, batch_size=38, shuffle=False)

    # models
    config = MyModelConfig(
        train_test_data.vocab_size_dic,
        HEADS,
        q_num=q_num,
        short_time=short_time,
        split_range=split_range,
        total_behavior_num=total_behavior_num,
        device_id=device_id,
        compressed=compressed,
        layers1=layers1,
        layers2=layers2,
        diff1=diff1,
        diff2=diff2
    )
    device = config.device

    if model_name == "myModel":
        model = MyModel(config)
    elif model_name == "myModelWithoutLayer2":
        model = MyModelWithoutLayer2(config)
    else:
        model = None

    model = model.to(device)

    checkpoint = torch.load(checkpoints_path, map_location=device, weights_only=False)
    model.load_state_dict(checkpoint['model_state_dict'])
    model.eval()
    print("模型数据准备完毕")

    return model, loader, config



def save_us_ts_MyModel(model_name,model, loader, config):

    model.eval()
    print("------保存压缩后的用户行为序列和兴趣向量------")
    device = config.device
    with torch.no_grad():

        for target_item, behaviors, global_interest_tokens, target,user_id in loader:
            print(user_id)
            target_item = target_item.to(device)
            behaviors = behaviors.to(device)
            global_interest_tokens = global_interest_tokens.to(device)

            behavior_embed,interest_embed,target_embed,behaviorsNumList = model.get_us_item(behaviors, global_interest_tokens,target_item)

            behavior_embed = behavior_embed.float().to(device) # b(1 or 38),behaviorNumList.size(),80
            interest_embed = interest_embed.float().to(device) # b(1 or 38),q_num,80
            target_embed = target_embed.float().to(device)  # b(1,38),1,80

            # 对于用户兴趣和用户的embedding进行选取
            behavior_embed = behavior_embed[0:1,:,:]  # 1 behaviorNumList.size() 80
            interest_embed = interest_embed[0:1,:,:]  # 1 query_num 80

            compressed_early_ts = behavior_embed[:, :behaviorsNumList[2], :] # 1,1,80
            compressed_middle_ts = behavior_embed[:, behaviorsNumList[2]:behaviorsNumList[2]+behaviorsNumList[1],:]  # 1,3,80
            latest_ts_other = behavior_embed[:, behaviorsNumList[2]+behaviorsNumList[1]:-short_time, :]  # 1,68,80
            latest_ts_short_time = behavior_embed[:, -short_time:, :]  # 1,20,80

            interest_us = interest_embed # 1, q_num,80

            # reshape
            interest_tensor = interest_us.reshape(-1,interest_us.shape[-1]) # query_num 80
            latest_other_tensor = latest_ts_other.reshape(-1, latest_ts_other.shape[-1]) #68 80
            latest_short_time_tensor = latest_ts_short_time.reshape(-1, latest_ts_other.shape[-1]) #20 80
            middle_tensor = compressed_middle_ts.reshape(-1, compressed_middle_ts.shape[-1]) #3,80
            early_tensor = compressed_early_ts.reshape(-1, compressed_early_ts.shape[-1]) #1,80
            target_tensor = target_embed.squeeze_(1)  # b(38,1) 80


            # 先转回 CPU 并转成 numpy
            interest_emb = interest_tensor.cpu().detach().numpy()
            latest_other_emb = latest_other_tensor.cpu().detach().numpy()
            latest_short_time_emb = latest_short_time_tensor.cpu().detach().numpy()
            middle_emb = middle_tensor.cpu().detach().numpy()
            early_emb = early_tensor.cpu().detach().numpy()
            item_emb = target_tensor.cpu().detach().numpy()


            # 拼接两个 tensor，一起降维
            all_emb = np.concatenate([item_emb, interest_emb,latest_other_emb,latest_short_time_emb,middle_emb,early_emb], axis=0)
            tsne = TSNE(n_components=2, perplexity=5, random_state=42, n_iter=1000)
            emb_2d = tsne.fit_transform(all_emb)  # shape: [N_total, 2]

            # 拆分回来
            item_2d = emb_2d[:len(item_emb)]
            interest_2d = emb_2d[len(item_emb):len(item_emb)+len(interest_emb)]
            latest_other_2d = emb_2d[len(item_emb)+len(interest_emb) : len(item_emb)+len(interest_emb)+len(latest_other_emb)]
            latest_short_time_2d = emb_2d[len(item_emb)+len(interest_emb)+len(latest_other_emb) : len(item_emb)+len(interest_emb)+len(latest_other_emb)+len(latest_short_time_emb)]
            middle_2d = emb_2d[len(item_emb)+len(interest_emb)++len(latest_other_emb)+len(latest_short_time_emb) : len(item_emb)+len(interest_emb)++len(latest_other_emb)+len(latest_short_time_emb)+len(middle_emb)]
            early_2d = emb_2d[len(item_emb)+len(interest_emb)++len(latest_other_emb)+len(latest_short_time_emb)+len(middle_emb):]

            # 绘图
            plt.figure(figsize=(8, 6))
            t_colors1 = ''
            t_colors2 = ''
            b_colors1_1 = '#aed1e7'
            b_colors1_2 = '#BEE5FC'
            b_colors2 = '#5ba3d1'
            b_colors3 = '#1c6bb0'
            for i in range(len(item_2d)):
                if target[i] == 1:
                    plt.scatter(item_2d[i, 0], item_2d[i, 1], c='green', label='Positive Item' if i == 0 else "", s=40)
                else:
                    plt.scatter(item_2d[i, 0], item_2d[i, 1], c='gray', label='Negative Item' if i == 0 else "", s=40)
            plt.scatter(interest_2d[:, 0], interest_2d[:, 1], c='red', label='Interest Tokens', marker='^', s=70)

            plt.scatter(latest_other_2d[:,0],latest_other_2d[:,1],c=b_colors1_1,label='Latest Tokens', marker='s', s=40)
            plt.scatter(latest_short_time_2d[:,0],latest_short_time_2d[:,1],c=b_colors1_2,label='Latest short_time Tokens', marker='s', s=40)

            plt.scatter(middle_2d[:,0],middle_2d[:,1],c=b_colors2,label='Middle Tokens', marker='s',  s=40)
            plt.scatter(early_2d[:,0],early_2d[:,1],c=b_colors3,label='Early Tokens',  marker='s',s=40)
            plt.title(f"t-SNE visualization of Items,Interest and Behavior Vectors of {user_id[0]}th user")
            plt.legend()
            plt.grid(True)
            plt.show()
            break

if __name__ == '__main__':
    model, loader, config = prepare_model_data(model_name, checkpoints_path="./checkpoints/myModel_7_final.pth", chose_loader='train')
    save_us_ts_MyModel(model_name,model, loader, config)
